use std::{io::{Read, Write, Seek, BufWriter, BufReader, SeekFrom}, time::Instant, fs::File};

use crate::{
    algo::{
        lzw,
        fano,
        huff,
        rle,
        bwt
    }, 
    common::{
        symcodec::SymbolCodec, 
        symbol::SymWithFreqU32
    }
};

fn universal_codec_test<C: SymbolCodec>(
    file_i: &'static str,
    codec_name: &'static str,
    codec_setup: impl FnOnce() -> C,
    encode_symbol_read_limit: u32
) {
    let file_o = format!("test-files/{codec_name}.eout");
    let file_o_d = format!("test-files/{codec_name}.dout");

    println!("Test for codec {codec_name}: ");
    
    let t0 = Instant::now();

    let codec = codec_setup();

    let d0 = t0.elapsed();

    println!("Codec setup takes: {:?}.", d0);

    let mut e_reader = BufReader::with_capacity(512 * 1024, File::open(file_i).unwrap());
    let mut e_writer = BufWriter::with_capacity(512 * 1024, File::create(&file_o).unwrap());
    codec.encode(&mut e_reader,&mut e_writer, encode_symbol_read_limit).unwrap();
    e_writer.flush().unwrap();

    let d1 = t0.elapsed();

    let original_size = e_reader.seek(SeekFrom::End(0)).unwrap();
    let encoded_size = e_writer.seek(SeekFrom::End(0)).unwrap();
    println!("Encoding (including IOs) takes {:?}, compression ratio = {}.", d1 - d0, original_size as f64 / encoded_size as f64);

    let d2 = t0.elapsed();

    let mut d_reader = BufReader::with_capacity(512 * 1024, File::open(&file_o).unwrap());
    let mut d_writer = BufWriter::with_capacity(512 * 1024, File::create(&file_o_d).unwrap());
    C::decode(&mut d_reader, &mut d_writer).unwrap();
    d_writer.flush().unwrap();

    let d3 = t0.elapsed();
    println!("Decoding (including IOs) takes {:?}.", d3 - d2);
    println!("Total times {:?}.", d3);
}

#[test]
fn fanocodec_u8symbol_test() {
    let d_full = vec![
        SymWithFreqU32 {sym: 65u8, freq: 31311141}, 
        SymWithFreqU32 {sym: 71, freq: 30947241}, 
        SymWithFreqU32 {sym: 67, freq: 30860341}, 
        SymWithFreqU32 {sym: 84, freq: 30718951}, 
        SymWithFreqU32 {sym: 48, freq: 713456}, 
        SymWithFreqU32 {sym: 49, freq: 600269}, 
        SymWithFreqU32 {sym: 50, freq: 434548}, 
        SymWithFreqU32 {sym: 95, freq: 294786}, 
        SymWithFreqU32 {sym: 115, freq: 245655}, 
        SymWithFreqU32 {sym: 54, freq: 219961}, 
        SymWithFreqU32 {sym: 51, freq: 210450}, 
        SymWithFreqU32 {sym: 56, freq: 171066}, 
        SymWithFreqU32 {sym: 53, freq: 162997}, 
        SymWithFreqU32 {sym: 99, freq: 147393}, 
        SymWithFreqU32 {sym: 55, freq: 131947}, 
        SymWithFreqU32 {sym: 10, freq: 98262}, 
        SymWithFreqU32 {sym: 47, freq: 98262}, 
        SymWithFreqU32 {sym: 110, freq: 98262}, 
        SymWithFreqU32 {sym: 52, freq: 95510}, 
        SymWithFreqU32 {sym: 57, freq: 80704}, 
        SymWithFreqU32 {sym: 62, freq: 49131}, 
        SymWithFreqU32 {sym: 101, freq: 49131}, 
        SymWithFreqU32 {sym: 109, freq: 49131}, 
        SymWithFreqU32 {sym: 111, freq: 49131}, 
        SymWithFreqU32 {sym: 112, freq: 49131}, 
        SymWithFreqU32 {sym: 117, freq: 49131}, 
        SymWithFreqU32 {sym: 78, freq: 25}
    ];
    universal_codec_test(
        "test-files/dataset.fasta",
        "fano",
        || {
            fano::FanoCodec::build(d_full, false)
        },
        u32::MAX
    );
}

#[test]
fn huffcodec_u8symbol_test() {
    let d_full = vec![
        SymWithFreqU32 {sym: 65u8, freq: 31311141}, 
        SymWithFreqU32 {sym: 71, freq: 30947241}, 
        SymWithFreqU32 {sym: 67, freq: 30860341}, 
        SymWithFreqU32 {sym: 84, freq: 30718951}, 
        SymWithFreqU32 {sym: 48, freq: 713456}, 
        SymWithFreqU32 {sym: 49, freq: 600269}, 
        SymWithFreqU32 {sym: 50, freq: 434548}, 
        SymWithFreqU32 {sym: 95, freq: 294786}, 
        SymWithFreqU32 {sym: 115, freq: 245655}, 
        SymWithFreqU32 {sym: 54, freq: 219961}, 
        SymWithFreqU32 {sym: 51, freq: 210450}, 
        SymWithFreqU32 {sym: 56, freq: 171066}, 
        SymWithFreqU32 {sym: 53, freq: 162997}, 
        SymWithFreqU32 {sym: 99, freq: 147393}, 
        SymWithFreqU32 {sym: 55, freq: 131947}, 
        SymWithFreqU32 {sym: 10, freq: 98262}, 
        SymWithFreqU32 {sym: 47, freq: 98262}, 
        SymWithFreqU32 {sym: 110, freq: 98262}, 
        SymWithFreqU32 {sym: 52, freq: 95510}, 
        SymWithFreqU32 {sym: 57, freq: 80704}, 
        SymWithFreqU32 {sym: 62, freq: 49131}, 
        SymWithFreqU32 {sym: 101, freq: 49131}, 
        SymWithFreqU32 {sym: 109, freq: 49131}, 
        SymWithFreqU32 {sym: 111, freq: 49131}, 
        SymWithFreqU32 {sym: 112, freq: 49131}, 
        SymWithFreqU32 {sym: 117, freq: 49131}, 
        SymWithFreqU32 {sym: 78, freq: 25}
    ];
    universal_codec_test(
        "test-files/dataset.fasta",
        "huff",
        || {
            huff::HuffmanCodec::build_from(d_full, false)
        },
        u32::MAX
    );
}

#[test]
fn lzw_u8symbol_test() {
    universal_codec_test(
        "test-files/dataset.fasta",
        "LZW", 
        || {
            lzw::LZWCodec::<u8>::build_for_full_nbits_symbol(12)
        }, 
        u32::MAX
    )
}

#[test]
fn rle_u8symbol_test() {
    universal_codec_test(
        "test-files/bwt-test-out-trans-all",
        "RLE", 
        || {
            rle::RLECodec::<u8>::new()
        }, 
        u32::MAX
    )
}

#[test]
fn bwt_test() {
    let file_i = "test-files/dataset_even_lines";
    let file_o_t = "test-files/bwt.tout";
    let file_o_c = "test-files/bwt.aux.c.tout";
    let file_o_occ = "test-files/bwt.aux.occ.tout";
    let file_o_dt = "test-files/bwt.rtout";

    let t = Instant::now();

    let mut f_orig_r = File::open(file_i).unwrap();
    let mut f_trans_w = BufWriter::with_capacity(128 * 1024, File::create(file_o_t).unwrap());
    let mut f_c_w = BufWriter::with_capacity(128 * 1024, File::create(file_o_c).unwrap());
    let mut f_occ_w = BufWriter::with_capacity(128 * 1024, File::create(file_o_occ).unwrap());

    let len = f_orig_r.seek(SeekFrom::End(0)).unwrap() as i32;
    f_orig_r.rewind().unwrap();
    println!("[bwt-test] File length: {}", len);

    bwt::transform(len, &mut f_orig_r, &mut f_trans_w, &mut f_c_w, &mut f_occ_w).unwrap();
    f_trans_w.flush().unwrap();
    f_c_w.flush().unwrap();
    f_occ_w.flush().unwrap();

    let d1 = t.elapsed();
    println!("[bwt-test] Encode time (including IOs): {:?}", d1);

    let mut f_trans_r = File::open(file_o_t).unwrap();
    let mut f_c_r = File::open(file_o_c).unwrap();
    let mut f_occ_r = File::open(file_o_occ).unwrap();
    let mut f_detrans_w = BufWriter::with_capacity(128 * 1024, File::create(file_o_dt).unwrap());

    bwt::reverse_transform(len, & mut f_trans_r, &mut f_c_r, &mut f_occ_r, &mut f_detrans_w, false);
    f_detrans_w.flush().unwrap();

    let d2 = t.elapsed();
    println!("[bwt-test] Decode time (including IOs): {:?}, total time: {:?}", d2 - d1, d2);
}
